# 用于给tensor组batch
import torch
import numpy as np

class groupByBatch:
    def __init__(self, batch_size):
        self.batch_size = int(batch_size)
        self.datas = []
        self.cnt = 0
        self.is_end = False
    
    # [Tensor,Tensor,...] -> [[],[],...,[Batch_Tensor],...]
    # DONE: 重写逻辑使其支持上述逻辑转换
    def process(self, data):
        res = []
        for i in data:
            self.datas.append(i.copy())
            self.cnt += 1
            if self.cnt >= self.batch_size: # 正好凑够了batch_size
                temp_res = np.array(self.datas)
                self.cnt = 0
                self.datas = []
                res.append([temp_res])
            else: # 还未凑够
                res.append([])
        return res

    def finish(self):
        res = []
        if self.cnt > 1:
            temp_res = np.array(self.datas)
            self.cnt = 0
            self.datas = []
            res.append([temp_res])
        else:
            res.append([])
        return res
            